import cv2
import logging
import coloredlogs
import numpy as np
import os
import sys
import yaml

# 配置日志
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')
# 创建一个logger实例
logging = logging.getLogger(__name__)

# 安装coloredlogs并自定义样式
coloredlogs.install(
    level='INFO',
    logger=logging,
    fmt='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    field_styles={
        'asctime': {'color': 'green'},
        'hostname': {'color': 'magenta'},
        'levelname': {'color': 'black', 'bold': True},
        'name': {'color': 'blue'},
    },
    level_styles={
        'debug': {'color': 'white'},
        'info': {'color': 'cyan'},
        'warning': {'color': 'yellow'},
        'error': {'color': 'red'},
        'critical': {'background': 'red'}
    }
)

class cfg_parser:
    def __init__(self, cfg_path):
        self.rec_size = 64
        self.pad = 0
        self.cfg_path = cfg_path
        
        self.parser()
        self.init_polygon()
        self.init_rectangle()
    
    def parser(self):
        with open(self.cfg_path, 'r') as f:
            self.cfg_data = yaml.safe_load(f)
            
            # 获取模板数量
            self.num_poly = [name[:7] for name in self.cfg_data].count("Polygon")
            self.num_rect = [name[:9] for name in self.cfg_data].count("Rectangle")
    def get_perspective_transform(self, roi):
        # 计算透视变换矩阵
        pts_dst = np.float32([[0, 0], [self.rec_size, 0], [
                                self.rec_size, self.rec_size], [0, self.rec_size]])  # 目标图像的四个顶点
        pts_roi = np.float32(roi)
        return (cv2.getPerspectiveTransform(pts_roi, pts_dst), pts_roi.reshape((-1, 1, 2)).astype(np.int32))
    
    def convert_axis(self, pts):
        # 坐标越界修复
        (xmin, ymin), (xmax, ymax) = pts
        xmin = xmin + self.pad if xmin < self.pad else xmin
        ymin = ymin + self.pad if ymin < self.pad else ymin
        # 构造ROI区域和边界框
        roi = [[xmin, ymax], [xmax, ymax], [xmax, ymin], [xmin, ymin]]
        box = [[xmin, ymin], [xmax, ymax]]
        return np.float32(roi).reshape((-1, 1, 2)).astype(np.int32), box
    def init_rectangle(self):
        # 定义矩形字典
        self.rect_dict = {key:None for key in [f"rect{num}" for num in range(self.num_rect)]}
        for i in range(self.num_rect):
            self.rect_dict[f'rect{i}'] = self.cfg_data[f'Rectangle{i}']
    
    def init_polygon(self):
         # 定义第一级和第二级的键
        first_level_keys = [
            f"poly{num}" for num in range(self.num_poly)]
        second_level_keys = ["M", "Pts"]
        
        self.poly_dict = {key: {
            sub_key: None for sub_key in second_level_keys} for key in first_level_keys}
        
        for i in range(self.num_poly):
            self.poly_dict[f"poly{i}"]["M"], self.poly_dict[f"poly{i}"]["Pts"] = self.get_perspective_transform(
                self.cfg_data[f"Polygon{i}"])
    
    
# 全局变量
polygon_points = []  # 存储多边形的顶点
rectangle_points = [] # 存储矩形的顶点
draw_count_poly = 0 # 用于记录当前绘制的多边形数量
draw_count_rect = 0 # 用于记录当前绘制的矩形数量
frame = None  # 当前帧
def template_image_save(video_path, save_path):
    global frame_count
     # 判断文件夹是否存在
    if not os.path.exists(save_path):
        logging.warning(f"The path does not exist. Create: {save_path}")
        try:
            os.makedirs(save_path)
        except OSError as e:
            logging.error(f"Failed to create path: {save_path}, Error: {e}")
            return
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)

    # 检查视频是否成功打开
    if not cap.isOpened():
        print("Error: Could not open video.")
        return

    paused = False  # 控制视频播放状态
    frame_count = 0  # 用于保存图片的计数器
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 获取总帧数
    # 绘制多边形的鼠标回调函数
    def draw_polygon_and_rectangle(event, x, y, flags, param):
        global draw_count_poly, draw_count_rect
        global polygon_points, rectangle_points
        global frame_count
        global save_path
        draw_frame = frame.copy()
        if event == cv2.EVENT_LBUTTONDOWN:
            polygon_points.append((x, y))
            cv2.circle(draw_frame, (x, y), 2, (0, 0, 255), -1)
        elif event == cv2.EVENT_LBUTTONUP:
            num = len(polygon_points)
            if num < 4:
                for i in range(num - 1):
                    cv2.line(draw_frame, polygon_points[i], polygon_points[i+1], (0, 255, 0), 2)
                cv2.imshow('Frame', draw_frame)
            else:
                cv2.polylines(draw_frame, [np.array(polygon_points)], isClosed=True, color=(0, 255, 0), thickness=2)
                cv2.imshow('Frame', draw_frame)           
                logging.info(f"Poly{draw_count_poly}保存结果:{polygon_points}")
                save_result_to_file(polygon_points, save_path)
                polygon_points = None
                polygon_points = []
                draw_count_poly += 1  
             
        # 绘制矩形坐标
        elif event == cv2.EVENT_MBUTTONDOWN:
            rectangle_points.append([x, y])
            cv2.circle(draw_frame, (x, y), 2, (0, 0, 255), -1)
            
        elif event == cv2.EVENT_MBUTTONUP:
            if len(rectangle_points) > 1:
                cv2.rectangle(draw_frame, rectangle_points[0], rectangle_points[1], (0, 255, 0), 2)
                cv2.imshow('Frame', draw_frame)
                logging.info(f"Rect{draw_count_rect}保存结果:{rectangle_points}")
                save_result_to_file(rectangle_points,save_path)
                rectangle_points = []
                draw_count_rect += 1
        
        # 鼠标左键保存落料图片
        if event == cv2.EVENT_RBUTTONDOWN:
            cfg_data = cfg_parser(os.path.join(save_path, 'roi.yaml'))
            # 判断文件是否存在
            if not os.path.exists(save_path):
                logging.warning(f"The path does not exist. Create:{save_path}")
                try:
                    os.makedirs(save_path)
                except OSError as e:
                    logging.error(f"Failed to create path:{save_path}, Error: {e}")
            # 落料检测区域校正图像  
            poly_img_list = np.array([cv2.warpPerspective(new_frame, cfg_data.poly_dict[f'poly{_num}']['M'], (cfg_data.rec_size, cfg_data.rec_size)) for _num in range(cfg_data.num_poly)])
            for i, img in enumerate(poly_img_list):
                cv2.imwrite(os.path.join(save_path, f"poly{i+1}.jpg"), img)
                logging.info(f'多边形区域{i}保存路径为:{os.path.join(save_path, f"poly{i+1}.jpg")}')
            # 保存定位区域模板图像
            for _num in range(cfg_data.num_rect):
                roi = cfg_data.rect_dict[f'rect{_num}']
                _, boxes = cfg_data.convert_axis(roi)
                img_save = new_frame[boxes[0][1]:boxes[1][1], boxes[0][0]:boxes[1][0]]
                cv2.imwrite(os.path.join(save_path, f'rect{_num}.jpg'), img_save)
                logging.info(f'矩形{_num}保存路径为:{os.path.join(save_path, f"rect{_num}.jpg")}')    
            frame_count += 1
    # 保存多边形坐标到文件
    def save_result_to_file(points, save_path):
         # 判断文件夹是否存在
        if not os.path.exists(save_path):
            logging.warning(f"The path does not exist. Create: {save_path}")
            try:
                os.makedirs(save_path)
            except OSError as e:
                logging.error(f"Failed to create path: {save_path}, Error: {e}")
                return

        file_path = os.path.join(save_path, 'roi.yaml')
        txt_path = os.path.join(save_path, 'roi.txt')
        
        if len(points) == 2:
            with open(txt_path, 'a') as f:
                f.write(f"Rectangle{draw_count_rect}:\n")
                f.write(f"[{rectangle_points[0]}, {rectangle_points[1]}]\n")
             # 准备要保存的数据
            save_data = {
                f"Rectangle{draw_count_rect}": [
                    [points[0][0], points[0][1]], [points[1][0], points[1][1]]
                ]
            }
        elif len(points) == 4:
            with open(txt_path, 'a') as f:
                f.write(f"Polygon{draw_count_poly}:\n")
                f.write(f"[[{points[0][0]},{points[0][1]}], [{points[1][0]},{points[1][1]}], [{points[2][0]},{points[2][1]}], [{points[3][0]},{points[3][1]}]]\n")
            # 准备要保存的数据
            save_data = {
                f"Polygon{draw_count_poly}": [
                    [points[0][0], points[0][1]], [points[1][0], points[1][1]], [points[2][0], points[2][1]], [points[3][0], points[3][1]]
                    ]
            }

        # 将数据写入 YAML 文件
        with open(file_path, 'a') as f:
            yaml.dump(save_data, f, allow_unicode=True, default_flow_style=False)

    # 设置鼠标回调
    cv2.namedWindow('Frame', cv2.WINDOW_NORMAL)
    cv2.setMouseCallback('Frame', draw_polygon_and_rectangle)

    while True:
        if not paused:
            ret, new_frame = cap.read()
            if not ret:
                print("Video has ended or failed to read frame.")
                break
            frame = new_frame

        # 显示帧
        if frame is not None:
            cv2.imshow('Frame', frame)

        # 更新进度条位置
        current_frame = int(cap.get(cv2.CAP_PROP_POS_FRAMES))

        # 捕获键盘输入
        key = cv2.waitKey(30) & 0xFF

        # 按下空格键暂停/继续
        if key == ord(' '):
            paused = not paused

        # 按下 q 键退出
        if key == ord('q'):
            break

    # 释放资源并关闭窗口
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    '''
    使用方法:空格暂停,左键点击绘制角点（绘制四个角点,按照左上角顺时针绘制）中键绘制矩形框（绘制两个点，左上角和右下角），鼠标右键保存结果，Q退出
    '''

    # 保存路径
    save_path = './temp'
    # 视频路径
    video_path = './type4.mov'
    #  绘图标识
    template_image_save(video_path, save_path)
    
    
    